/*
 *  Copyright 2020-2025 the original author or authors.
 *  You cannot use this file unless authorized by the author.
 */

package org.ipig.computing.hadoop.yarn;

import com.alibaba.fastjson.JSONObject;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang.StringUtils;
import org.apache.hadoop.yarn.api.records.ApplicationReport;
import org.apache.hadoop.yarn.api.records.YarnApplicationState;
import org.apache.hadoop.yarn.client.api.YarnClient;
import org.apache.hadoop.yarn.conf.YarnConfiguration;
import org.ipig.commons.helper.HttpHelper;
import org.ipig.constants.ResultStatus;
import org.ipig.constants.SymbolCnst;

import java.util.EnumSet;
import java.util.List;

/**
 * YarnHelper
 *
 * @author <a href="mailto:comchnts@163.com">chinats</a>
 * @since 1.0
 */
@Slf4j
public class YarnHelper {
    /**
     * checkAppRunning
     *
     * @param appName
     * @param haActiveHost
     * @return
     */
    public static ResultStatus checkAppRunning(final String appName, final String haActiveHost) {
        ResultStatus resultStatus = ResultStatus.SUCCESS;
        String rsAddress = new StringBuilder().append(haActiveHost).append(SymbolCnst.COLON).append("8032").toString();
        String rsSchedulerAddress = new StringBuilder().append(haActiveHost).append(SymbolCnst.COLON).append("8030").toString();
        String rsTrackerAddress = new StringBuilder().append(haActiveHost).append(SymbolCnst.COLON).append("8030").toString();
        String appId = "";
        YarnConfiguration yarnConf = new YarnConfiguration();
        yarnConf.setStrings("yarn.resourcemanager.scheduler.address", rsSchedulerAddress);
        yarnConf.setStrings("yarn.resourcemanager.resource-tracker.address", rsTrackerAddress);
        yarnConf.set("yarn.resourcemanager.address", rsAddress);
        EnumSet<YarnApplicationState> applicationStates = EnumSet.of(YarnApplicationState.RUNNING, YarnApplicationState.ACCEPTED);
        try {
            List<ApplicationReport> applications = YarnHelper.getApplications(yarnConf, applicationStates);
            if (applications != null && !applications.isEmpty()) {
                for (ApplicationReport ar : applications) {
                    if (StringUtils.equals(ar.getName(), appName)) {
                        appId = ar.getApplicationId().toString();
                        break;
                    }
                }
            }
            if (StringUtils.isNotBlank(appId)) {
                resultStatus = ResultStatus.FAILURE;
                log.warn("应用【{}】，ID【{}】正在运行，不再向集群提交。", appName, appId);
            }
        } catch (Exception ex) {
            ex.printStackTrace();
        }
        return resultStatus;
    }

    /**
     * getHaActiveHost
     * @param hostArgs
     * @return
     */
    public static String getHaActiveHost(String hostArgs) {
        String host = "";
        String[] hosts = StringUtils.split(hostArgs, SymbolCnst.COMMA);
        if (hosts != null && hosts.length > 0) {
            for (int i = 0; i < hosts.length; i++) {
                String hostAndPort = hosts[i];
                StringBuilder clusterInfoUrl = new StringBuilder();
                clusterInfoUrl.append("http://").append(hostAndPort).append("/ws/v1/cluster/info");
                String result = HttpHelper.doGet(clusterInfoUrl.toString());
                log.debug(result);
                if (!StringUtils.isBlank(result)) {
                    JSONObject jsonObject = JSONObject.parseObject(result);
                    if (jsonObject != null) {
                        JSONObject clusterInfo = (JSONObject) jsonObject.get("clusterInfo");
                        if (clusterInfo != null) {
                            if (StringUtils.equals(clusterInfo.getString("haState"), "ACTIVE")) {
                                host = StringUtils.substring(hostAndPort, 0, StringUtils.indexOf(hostAndPort, SymbolCnst.COLON));
                                break;
                            }
                        }
                    }
                }
            }
        }
        return host;
    }

    /**
     * getApplications
     *
     * @param conf
     * @param applicationStates
     * @return
     */
    public static List<ApplicationReport> getApplications(YarnConfiguration conf, EnumSet<YarnApplicationState> applicationStates) {
        List<ApplicationReport> applications = null;
        YarnClient yarnClient = null;
        try {
            // YarnConfiguration继承自Configuration
            yarnClient = YarnClient.createYarnClient();
            yarnClient.init(conf);
            yarnClient.start();
            applications = yarnClient.getApplications(applicationStates);
        } catch (Exception ex) {
            ex.printStackTrace();
        } finally {
            if (yarnClient != null) {
                yarnClient.stop();
            }
        }
        return applications;
    }
}

